from celery import Celery
import redis
import json
from services import HSKService
celery_app = Celery('tasks', broker='redis://localhost:6379/0', backend='redis://localhost:6379/1')
r = redis.Redis(host='localhost', port=6379, db=2)

def send_to_socketio(sid, step, result):
    # 将消息写入 Redis，由 SocketIO 服务轮询或通过通道发送
    message = json.dumps({'sid': sid, 'step': step, 'result': result})
    r.publish('steps', message)

@celery_app.task
def run_steps(topic, level, sid):
    print('执行任务的时候',sid)

    services = HSKService()
    theme_labels = services.get_theme_labels(topic)
    send_to_socketio(sid,1, {"theme_labels": theme_labels})
    hsk_words = services.get_words_sample(topic,level)
    combined_scores = {**hsk_words["scores"], **hsk_words["supplement_scores"]}
    # 按分数降序排序
    sorted_words = sorted(combined_scores.items(), key=lambda x: x[1], reverse=True)
    # 取前10(可以根据需要调整)个语义最相关的词汇
    send_to_socketio(sid, 2, {"hsk_words":hsk_words})
    top_words = [word for word, score in sorted_words[:10]]  # 提取前10个单词（去掉分数）
    # 流式输出，在这个函数中已经实现了向前端推送了，也就是第三步的推送generate_article_stream
    print('执行异步恩物之前的时候', sid)
    passage = services.generate_article_stream(topic,sid, level, top_words)
    send_to_socketio(sid, step=4, result='done')  # 只发这段，不是整篇
    # 非流式输出
    # passage = services.generate_article(topic, level, top_words)
    # send_to_socketio(sid, 3, {"article": passage})
    # 断开webSocket的标识

    return "done"
